package org.spiderflow.mongodb.utils;

import com.mongodb.client.MongoDatabase;
import com.mongodb.client.gridfs.GridFSBucket;
import com.mongodb.client.gridfs.GridFSBuckets;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import org.apache.commons.lang3.StringUtils;
import org.spiderflow.mongodb.model.MongoDataSource;

import com.mongodb.MongoClient;
import com.mongodb.MongoClientOptions;
import com.mongodb.MongoClientOptions.Builder;
import com.mongodb.MongoCredential;
import com.mongodb.ServerAddress;
import org.springframework.data.convert.CustomConversions;
import org.springframework.data.mongodb.MongoDbFactory;
import org.springframework.data.mongodb.config.MongoConfigurationSupport;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.SimpleMongoDbFactory;
import org.springframework.data.mongodb.core.convert.DbRefResolver;
import org.springframework.data.mongodb.core.convert.DefaultDbRefResolver;
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
import org.springframework.data.mongodb.gridfs.GridFsTemplate;

public class MongoDBUtils {

  private static Map<String, GridFSBucket> cacheGridFSBucket = new HashMap<>();
  private static Map<String, MongoTemplate> cacheMongoTemplate = new HashMap<>();

  public static MongoClient createMongoClient(MongoDataSource dataSource) {
    ServerAddress address = new ServerAddress(dataSource.getHost(), dataSource.getPort());
    Builder options = new MongoClientOptions.Builder();
    options.connectTimeout(5000)
        .maxWaitTime(3000)
        .serverSelectionTimeout(3000);
    if (StringUtils.isNotBlank(dataSource.getUsername()) && StringUtils
        .isNotBlank(dataSource.getPassword())) {
      MongoCredential credential = MongoCredential
          .createScramSha1Credential(dataSource.getUsername(), dataSource.getDatabase(),
              dataSource.getPassword().toCharArray());
      return new MongoClient(address, credential, options.build());
    } else {
      return new MongoClient(address, options.build());
    }
  }

  //GridFSBucket用于打开下载流
  public static GridFSBucket getGridFSBucket(MongoClient mongoClient, String db) {
    GridFSBucket client = cacheGridFSBucket.get(db);
    if (null != client) {
      return client;
    }

    MongoDatabase mongoDatabase = mongoClient.getDatabase(db);
    GridFSBucket bucket = GridFSBuckets.create(mongoDatabase);
    cacheGridFSBucket.put(db, bucket);
    return bucket;
  }


  public static MongoTemplate getMongoTemplate(MongoClient mongoClient, String db) {
    MongoTemplate mongoTemplate = cacheMongoTemplate.get(db);
    if (null != mongoTemplate) {
      return mongoTemplate;
    }

    mongoTemplate = new MongoTemplate(mongoClient, db);
    cacheMongoTemplate.put(db, mongoTemplate);
    return mongoTemplate;
  }

  public static GridFsTemplate gridFsTemplate(MongoClient mongoClient, String db) throws Exception {
    return new GridFsTemplate(mongoDbFactory(mongoClient, db), mappingMongoConverter(mongoClient, db));
  }

  private static MongoDbFactory mongoDbFactory(MongoClient mongoClient, String db) throws Exception {
    return new SimpleMongoDbFactory(mongoClient, db);
  }

  private static MappingMongoConverter mappingMongoConverter(MongoClient mongoClient, String db) throws Exception {
    DbRefResolver dbRefResolver = new DefaultDbRefResolver(mongoDbFactory(mongoClient, db));
    MongoConfigurationSupport mongoConfigurationSupport = new MongoConfigurationSupport() {

      @Override
      protected String getDatabaseName() {
        return db;
      }
    };

    MappingMongoConverter converter = new MappingMongoConverter(dbRefResolver, mongoConfigurationSupport.mongoMappingContext());
    converter.setCustomConversions(customMyConversions());
    return converter;
  }

  private static CustomConversions customMyConversions() {
    return new MongoCustomConversions(Collections.emptyList());
  }
}
